import torch
import torch.nn as nn
from .util import init

"""
修改标准的PyTorch分布，使其与此代码库兼容。
"""

#
# 标准化分布接口
#

# 分类分布
class FixedCategorical(torch.distributions.Categorical):
    """
    固定的分类分布类，继承自PyTorch的Categorical分布
    """
    def sample(self):
        """
        采样一个动作并在最后一个维度添加一个维度
        """
        return super().sample().unsqueeze(-1)

    def log_probs(self, actions):
        """
        计算给定动作的对数概率
        
        参数:
            actions: 要计算对数概率的动作
            
        返回:
            动作的对数概率
        """
        return (
            super()
            .log_prob(actions.squeeze(-1))
            .view(actions.size(0), -1)
            .sum(-1)
            .unsqueeze(-1)
        )

    def mode(self):
        """
        返回概率最高的动作
        """
        return self.probs.argmax(dim=-1, keepdim=True)


# 正态分布
class FixedNormal(torch.distributions.Normal):
    """
    固定的正态分布类，继承自PyTorch的Normal分布
    """
    def log_probs(self, actions):
        """
        计算给定动作的对数概率
        
        参数:
            actions: 要计算对数概率的动作
            
        返回:
            动作的对数概率
        """
        return super().log_prob(actions).sum(-1, keepdim=True)

    def entrop(self):
        """
        计算分布的熵
        """
        return super.entropy().sum(-1)

    def mode(self):
        """
        返回分布的均值作为模式
        """
        return self.mean


# 伯努利分布
class FixedBernoulli(torch.distributions.Bernoulli):
    """
    固定的伯努利分布类，继承自PyTorch的Bernoulli分布
    """
    def log_probs(self, actions):
        """
        计算给定动作的对数概率
        
        参数:
            actions: 要计算对数概率的动作
            
        返回:
            动作的对数概率
        """
        return super.log_prob(actions).view(actions.size(0), -1).sum(-1).unsqueeze(-1)

    def entropy(self):
        """
        计算分布的熵
        """
        return super().entropy().sum(-1)

    def mode(self):
        """
        返回概率大于0.5的动作
        """
        return torch.gt(self.probs, 0.5).float()


class Categorical(nn.Module):
    """
    分类分布模块，用于离散动作空间
    """
    def __init__(self, num_inputs, num_outputs, use_orthogonal=True, gain=0.01):
        """
        初始化分类分布模块
        
        参数:
            num_inputs: 输入特征的维度
            num_outputs: 输出动作的维度
            use_orthogonal: 是否使用正交初始化
            gain: 初始化的增益
        """
        super(Categorical, self).__init__()
        init_method = [nn.init.xavier_uniform_, nn.init.orthogonal_][use_orthogonal]
        def init_(m): 
            return init(m, init_method, lambda x: nn.init.constant_(x, 0), gain)

        # 线性层，将输入特征映射到动作空间
        self.linear = init_(nn.Linear(num_inputs, num_outputs))

    def forward(self, x, available_actions=None):
        """
        前向传播
        
        参数:
            x: 输入特征
            available_actions: 可用动作的掩码
            
        返回:
            分类分布
        """
        x = self.linear(x)
        # 如果提供了可用动作掩码，将不可用动作的概率设为极小值
        if available_actions is not None:
            x[available_actions == 0] = -1e10
        return FixedCategorical(logits=x)


class DiagGaussian(nn.Module):
    """
    对角高斯分布模块，用于连续动作空间
    """
    def __init__(self, num_inputs, num_outputs, use_orthogonal=True, gain=0.01):
        """
        初始化对角高斯分布模块
        
        参数:
            num_inputs: 输入特征的维度
            num_outputs: 输出动作的维度
            use_orthogonal: 是否使用正交初始化
            gain: 初始化的增益
        """
        super(DiagGaussian, self).__init__()

        init_method = [nn.init.xavier_uniform_, nn.init.orthogonal_][use_orthogonal]
        def init_(m): 
            return init(m, init_method, lambda x: nn.init.constant_(x, 0), gain)

        # 均值网络
        self.fc_mean = init_(nn.Linear(num_inputs, num_outputs))
        # 标准差网络
        self.logstd = AddBias(torch.zeros(num_outputs))

    def forward(self, x):
        """
        前向传播
        
        参数:
            x: 输入特征
            
        返回:
            正态分布
        """
        # 计算动作均值
        action_mean = self.fc_mean(x)

        # 为KFAC实现的一个丑陋的hack
        zeros = torch.zeros(action_mean.size())
        if x.is_cuda:
            zeros = zeros.cuda()

        # 计算动作对数标准差
        action_logstd = self.logstd(zeros)
        return FixedNormal(action_mean, action_logstd.exp())


class Bernoulli(nn.Module):
    """
    伯努利分布模块，用于二元动作空间
    """
    def __init__(self, num_inputs, num_outputs, use_orthogonal=True, gain=0.01):
        """
        初始化伯努利分布模块
        
        参数:
            num_inputs: 输入特征的维度
            num_outputs: 输出动作的维度
            use_orthogonal: 是否使用正交初始化
            gain: 初始化的增益
        """
        super(Bernoulli, self).__init__()
        init_method = [nn.init.xavier_uniform_, nn.init.orthogonal_][use_orthogonal]
        def init_(m): 
            return init(m, init_method, lambda x: nn.init.constant_(x, 0), gain)
        
        # 线性层，将输入特征映射到动作空间
        self.linear = init_(nn.Linear(num_inputs, num_outputs))

    def forward(self, x):
        """
        前向传播
        
        参数:
            x: 输入特征
            
        返回:
            伯努利分布
        """
        x = self.linear(x)
        return FixedBernoulli(logits=x)

class AddBias(nn.Module):
    """
    添加偏置的模块
    """
    def __init__(self, bias):
        """
        初始化添加偏置的模块
        
        参数:
            bias: 要添加的偏置
        """
        super(AddBias, self).__init__()
        self._bias = nn.Parameter(bias.unsqueeze(1))

    def forward(self, x):
        """
        前向传播
        
        参数:
            x: 输入张量
            
        返回:
            添加了偏置的张量
        """
        if x.dim() == 2:
            bias = self._bias.t().view(1, -1)
        else:
            bias = self._bias.t().view(1, -1, 1, 1)

        return x + bias
